import json
import pandas as pd
import tkinter as tk
from tkinter import filedialog
import os

def select_file():
    """弹出文件选择窗口，返回选定的 JSONL 文件路径"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    file_path = filedialog.askopenfilename(title="选择 JSONL 数据文件", filetypes=[("JSONL 文件", "*.jsonl")])
    return file_path

def clean_and_select_features(file_path):
    """清洗 JSONL 数据并选择重要特征"""
    cleaned_data = []

    # 逐行读取 JSONL 文件
    with open(file_path, "r", encoding="utf-8") as file:
        for line in file:
            entry = json.loads(line.strip())  # 解析 JSON 对象
            new_entry = {key.strip(): value.strip() if isinstance(value, str) else value for key, value in entry.items()}
            cleaned_data.append(new_entry)

    # 转换为 Pandas DataFrame
    df = pd.DataFrame(cleaned_data)

    # 转换数据类型（去除空字符串并转换为数值类型）
    for col in df.columns:
        if col != 'Label':
            df[col] = pd.to_numeric(df[col], errors="coerce")

    # 删除全为 0 的特征
    df = df.loc[:, (df != 0).any(axis=0)]

    # 删除方差接近 0 的特征（低信息量特征），跳过 'Label' 列
    variance_threshold = 1e-5  # 方差阈值
    # 选择除 Label 之外的所有数值型列
    numeric_cols = [col for col in df.columns if col != 'Label']
    var_series = df[numeric_cols].var()
    selected_numeric_cols = var_series[var_series > variance_threshold].index.tolist()
    
    # 如果存在 'Label' 列，则把它加回去
    if 'Label' in df.columns:
        selected_cols = selected_numeric_cols + ['Label']
    else:
        selected_cols = selected_numeric_cols

    df = df[selected_cols]

    # 获取清洗后的 JSONL 文件路径
    cleaned_file_name = f"cleaned_{os.path.basename(file_path)}"
    cleaned_file_path = os.path.join(os.path.dirname(file_path), cleaned_file_name)

    # 逐行写入 JSONL 文件
    with open(cleaned_file_path, "w", encoding="utf-8") as cleaned_file:
        for record in df.to_dict(orient="records"):
            cleaned_file.write(json.dumps(record, ensure_ascii=False) + "\n")

    print(f"数据清洗完成，清洗后的文件保存为：{cleaned_file_path}")

# 运行数据清洗流程
if __name__ == "__main__":
    file_path = select_file()
    if file_path:
        clean_and_select_features(file_path)
    else:
        print("未选择文件，程序退出。")
